#pragma once
#include "copy.h"
#include <iostream>

/**
 * Wrapper around cublasSgemm() that extracts the arguments and performs the
 * translation necessary to go from copperhead data types to data types cuBLAS
 * is expecting. performs D = alpha*A*B + beta*C.
 *
 * @param _transA whether or not to transpose matrix A. 0 is no, 1 is yes
 * @param _transB whether or not to transpose matrix B.
 * @param alpha scaling factor for the Matrix A*B
 * @param A m by k Matrix to be scaled and multiplied
 * @param B k by n Matrix to be multiplied
 * @param beta scaling factor for Matrix C
 * @param C m by n Matrix to be scaled and added to A*B
 * @return D = alpha*A*B + beta*C in column major format.
 */
 
void sgemm
(
    int _transA, 
    int _transB, 
    float alpha, 
    Matrix A, 
    Matrix B, 
    float beta, 
    Matrix C,
    Matrix D
)
{
    bool is_col_A, is_col_B, is_square_A, is_square_B, opt_A, opt_B;
    int m, n, k, lda, ldb, ldc;
    char transA, transB;

    is_col_A = (A.stride == 1);
    is_col_B = (B.stride == 1);

    is_square_A = (A.length == A.sub.length);
    is_square_B = (B.length == B.sub.length);

    if(!is_col_A && !is_square_A && !_transA){ //we have to copy A
        std::cerr << "sgemm: WARNING: Passing in a non-square, row major matrix without transposing has incurred a copy for A" << std::endl;
        int size_A = A.length * A.sub.length * sizeof(float);
        float *A_data;
        cudaMalloc(&A_data, size_A);
        SubMatrix A_sub(A.sub.length, A.length, A_data, 0);
        Matrix A_copy(A.length, 1, A_sub);
        copy(A, A_copy);
        A = A_copy;
        is_col_A = 1;
    }

    if(!is_col_B && !is_square_B && !_transB){ //we have to copy B
        std::cerr << "sgemm: WARNING: Passing in a non-square, row major matrix without transposing has incurred a copy for B" << std::endl;
        int size_B = B.length * B.sub.length * sizeof(float);
        float *B_data;
        cudaMalloc(&B_data, size_B);
        SubMatrix B_sub(B.sub.length, B.length, B_data, 0);
        Matrix B_copy(B.length, 1, B_sub);
        copy(B, B_copy);
        B = B_copy;
        is_col_B = 1;
    }

    /* get our dimensions, by convention in copperhead
    the first nested depth length is number of rows */
    opt_A = !is_col_A && _transA;
    opt_B = !is_col_B && _transB;
    m = opt_A ? A.sub.length : A.length;
    k = opt_A ? A.length : A.sub.length;
    n = D.sub.length;

    if(m != D.length){ //dimension mismatch
        std::cerr << "sgemm: ERROR: Dimension mismatch in matrices A and C: (" 
            << m << " and " << D.length << " should be equal" << std::endl;
        return;
    }

    if(k != (opt_B || _transB ? B.sub.length : B.length)){ 
        std::cerr << "sgemm: ERROR: Dimension mismatch in matrices A and B: (" 
            << k << " and " << (opt_B ? B.sub.length : B.length) 
            << " should be equal" << std::endl;
        return;
    }

    if((_transB ? B.length : B.sub.length) != D.sub.length){ 
        std::cerr << "sgemm: ERROR: Dimension mismatch in matrices B and C: (" 
            << n << " and " << D.sub.length << " should be equal" << std::endl;
        return;
    }

    // calculate ld{a,b,c}
    lda = is_col_A ? A.sub.stride : A.stride;
    ldb = is_col_B ? B.sub.stride : B.stride;
    ldc = D.sub.stride;

    if(D.stride != 1){
        std::cerr << "sgemm: ERROR - D must be allocated column major" << std::endl;
        return;
    }
    if(_transA > 1 || _transB > 1){
        std::cerr << "sgemm: ERROR - invalid value for transA or transB, must be 0 or 1" << std::endl;
        return;
    }
    // copy C before feeding it to sgemm (also ensuring it's column major)
    copy(C, D);
    

    transA = (_transA ^ is_col_A) ? 'n' : 't';
    transB = (_transB ^ is_col_B) ? 'n' : 't';

    cublasSgemm(transA, transB, m, n, k, alpha, A.sub.data + A.sub.offset, lda, B.sub.data + B.sub.offset, ldb, beta, D.sub.data + D.sub.offset, ldc);
}
